{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of text in characters 1115394\n",
      "First Citizen:\n",
      "Before we proceed any further, hear me speak.\n",
      "\n",
      "All:\n",
      "Speak, speak.\n",
      "\n",
      "First Citizen:\n",
      "You\n"
     ]
    }
   ],
   "source": [
    "# 读入数据\n",
    "with open(\"./dataset.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "    text = f.read()\n",
    "text_len = len(text)\n",
    "# —————————————————————————————————————————————————————— #\n",
    "print(f\"length of text in characters {text_len}\")\n",
    "print(text[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "there are 65 characters totally\n",
      "they are \n",
      " !$&',-.3:;?ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz\n"
     ]
    }
   ],
   "source": [
    "chars = sorted(list(set(text))) # 提取所有不重复字符并放在一个列表中排序\n",
    "vocabs_len = len(chars)\n",
    "# —————————————————————————————————————————————————————— #\n",
    "print(f\"there are {vocabs_len} characters totally\")\n",
    "print(\"they are \"+\"\".join(chars))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[20, 43, 50, 50, 53, 1, 61, 53, 56, 50, 42, 2]\n",
      "Hello world!\n"
     ]
    }
   ],
   "source": [
    "# 这里我们对字符和标点符号进行编码\n",
    "# 编码器encoder将字符串翻译成整数列表\n",
    "# 解码器decoder将整数列表翻译成字符串\n",
    "stoi = {ch:idx for idx,ch in enumerate(chars)}\n",
    "itos = chars\n",
    "encoder = lambda s: [stoi[c] for c in s]\n",
    "decoder = lambda idxs: \"\".join([itos[idx] for idx in idxs])\n",
    "# —————————————————————————————————————————————————————— #\n",
    "print(encoder(\"Hello world!\"))\n",
    "print(decoder(encoder(\"Hello world!\")))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([18, 47, 56, 57, 58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 14, 43, 44,\n",
      "        53, 56, 43,  1, 61, 43,  1, 54, 56, 53, 41, 43, 43, 42,  1, 39, 52, 63,\n",
      "         1, 44, 59, 56, 58, 46, 43, 56,  6,  1, 46, 43, 39, 56,  1, 51, 43,  1,\n",
      "        57, 54, 43, 39, 49,  8,  0,  0, 13, 50, 50, 10,  0, 31, 54, 43, 39, 49,\n",
      "         6,  1, 57, 54, 43, 39, 49,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47,\n",
      "        58, 47, 64, 43, 52, 10,  0, 37, 53, 59,  1, 39, 56, 43,  1, 39, 50, 50,\n",
      "         1, 56, 43, 57, 53, 50, 60, 43, 42,  1, 56, 39, 58, 46, 43, 56,  1, 58,\n",
      "        53,  1, 42, 47, 43,  1, 58, 46, 39, 52,  1, 58, 53,  1, 44, 39, 51, 47,\n",
      "        57, 46, 12,  0,  0, 13, 50, 50, 10,  0, 30, 43, 57, 53, 50, 60, 43, 42,\n",
      "         8,  1, 56, 43, 57, 53, 50, 60, 43, 42,  8,  0,  0, 18, 47, 56, 57, 58,\n",
      "         1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 18, 47, 56, 57, 58,  6,  1, 63,\n",
      "        53, 59,  1, 49, 52, 53, 61,  1, 15, 39, 47, 59, 57,  1, 25, 39, 56, 41,\n",
      "        47, 59, 57,  1, 47, 57,  1, 41, 46, 47, 43, 44,  1, 43, 52, 43, 51, 63,\n",
      "         1, 58, 53,  1, 58, 46, 43,  1, 54, 43, 53, 54, 50, 43,  8,  0,  0, 13,\n",
      "        50, 50, 10,  0, 35, 43,  1, 49, 52, 53, 61,  5, 58,  6,  1, 61, 43,  1,\n",
      "        49, 52, 53, 61,  5, 58,  8,  0,  0, 18, 47, 56, 57, 58,  1, 15, 47, 58,\n",
      "        47, 64, 43, 52, 10,  0, 24, 43, 58,  1, 59, 57,  1, 49, 47, 50, 50,  1,\n",
      "        46, 47, 51,  6,  1, 39, 52, 42,  1, 61, 43,  5, 50, 50,  1, 46, 39, 60,\n",
      "        43,  1, 41, 53, 56, 52,  1, 39, 58,  1, 53, 59, 56,  1, 53, 61, 52,  1,\n",
      "        54, 56, 47, 41, 43,  8,  0, 21, 57,  5, 58,  1, 39,  1, 60, 43, 56, 42,\n",
      "        47, 41, 58, 12,  0,  0, 13, 50, 50, 10,  0, 26, 53,  1, 51, 53, 56, 43,\n",
      "         1, 58, 39, 50, 49, 47, 52, 45,  1, 53, 52,  5, 58, 11,  1, 50, 43, 58,\n",
      "         1, 47, 58,  1, 40, 43,  1, 42, 53, 52, 43, 10,  1, 39, 61, 39, 63,  6,\n",
      "         1, 39, 61, 39, 63,  2,  0,  0, 31, 43, 41, 53, 52, 42,  1, 15, 47, 58,\n",
      "        47, 64, 43, 52, 10,  0, 27, 52, 43,  1, 61, 53, 56, 42,  6,  1, 45, 53,\n",
      "        53, 42,  1, 41, 47, 58, 47, 64, 43, 52, 57,  8,  0,  0, 18, 47, 56, 57,\n",
      "        58,  1, 15, 47, 58, 47, 64, 43, 52, 10,  0, 35, 43,  1, 39, 56, 43,  1,\n",
      "        39, 41, 41, 53, 59, 52, 58, 43, 42,  1, 54, 53, 53, 56])\n"
     ]
    }
   ],
   "source": [
    "# 对整个文本进行编码并转化成long tensor\n",
    "import torch\n",
    "# data是全文的整数翻译\n",
    "data = torch.tensor(encoder(text), dtype=torch.long)\n",
    "# —————————————————————————————————————————————————————— #\n",
    "print(data[:500])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 把数据分成训练集和验证集，后期验证是否过拟合了\n",
    "n_90 = int(text_len*0.9)\n",
    "train_data = data[:n_90]\n",
    "val_data = data[n_90:]\n",
    "# —————————————————————————————————————————————————————— #"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "inputs shape torch.Size([4, 8])\n",
      "tensor([[53, 59,  6,  1, 58, 56, 47, 40],\n",
      "        [49, 43, 43, 54,  1, 47, 58,  1],\n",
      "        [13, 52, 45, 43, 50, 53,  8,  0],\n",
      "        [ 1, 39,  1, 46, 53, 59, 57, 43]])\n",
      "target shape torch.Size([4, 8])\n",
      "tensor([[59,  6,  1, 58, 56, 47, 40, 59],\n",
      "        [43, 43, 54,  1, 47, 58,  1, 58],\n",
      "        [52, 45, 43, 50, 53,  8,  0, 26],\n",
      "        [39,  1, 46, 53, 59, 57, 43,  0]])\n",
      "input tensor([53]) have target 59\n",
      "input tensor([53, 59]) have target 6\n",
      "input tensor([53, 59,  6]) have target 1\n",
      "input tensor([53, 59,  6,  1]) have target 58\n",
      "input tensor([53, 59,  6,  1, 58]) have target 56\n",
      "input tensor([53, 59,  6,  1, 58, 56]) have target 47\n",
      "input tensor([53, 59,  6,  1, 58, 56, 47]) have target 40\n",
      "input tensor([53, 59,  6,  1, 58, 56, 47, 40]) have target 59\n",
      "input tensor([49]) have target 43\n",
      "input tensor([49, 43]) have target 43\n",
      "input tensor([49, 43, 43]) have target 54\n",
      "input tensor([49, 43, 43, 54]) have target 1\n",
      "input tensor([49, 43, 43, 54,  1]) have target 47\n",
      "input tensor([49, 43, 43, 54,  1, 47]) have target 58\n",
      "input tensor([49, 43, 43, 54,  1, 47, 58]) have target 1\n",
      "input tensor([49, 43, 43, 54,  1, 47, 58,  1]) have target 58\n",
      "input tensor([13]) have target 52\n",
      "input tensor([13, 52]) have target 45\n",
      "input tensor([13, 52, 45]) have target 43\n",
      "input tensor([13, 52, 45, 43]) have target 50\n",
      "input tensor([13, 52, 45, 43, 50]) have target 53\n",
      "input tensor([13, 52, 45, 43, 50, 53]) have target 8\n",
      "input tensor([13, 52, 45, 43, 50, 53,  8]) have target 0\n",
      "input tensor([13, 52, 45, 43, 50, 53,  8,  0]) have target 26\n",
      "input tensor([1]) have target 39\n",
      "input tensor([ 1, 39]) have target 1\n",
      "input tensor([ 1, 39,  1]) have target 46\n",
      "input tensor([ 1, 39,  1, 46]) have target 53\n",
      "input tensor([ 1, 39,  1, 46, 53]) have target 59\n",
      "input tensor([ 1, 39,  1, 46, 53, 59]) have target 57\n",
      "input tensor([ 1, 39,  1, 46, 53, 59, 57]) have target 43\n",
      "input tensor([ 1, 39,  1, 46, 53, 59, 57, 43]) have target 0\n"
     ]
    }
   ],
   "source": [
    "# 获取批量，每批的大小是block_size\n",
    "'''\n",
    "x:012345\n",
    "y:123456\n",
    "block_y中的第i位对应着x前i位的预测\n",
    "'''\n",
    "torch.manual_seed(1337)\n",
    "batch_size = 4\n",
    "block_size = 8\n",
    "# 获取一批输入输出，y的位置对应x的target\n",
    "def get_batch(mode=\"train\"):\n",
    "    data = train_data if mode == \"train\" else val_data\n",
    "    idxs = torch.randint(0, n_90-1-block_size, (batch_size,)) # 生成一个起始索引\n",
    "    x = torch.stack([data[idx:idx+block_size] for idx in idxs]) # 根据索引生成批量数据(batch_size, block_size)\n",
    "    y = torch.stack([data[idx+1:idx+block_size+1] for idx in idxs]) # x往后偏移一位\n",
    "    return x, y\n",
    "# —————————————————————————————————————————————————————— #\n",
    "bx, by = get_batch()\n",
    "print(f\"inputs shape {bx.shape}\\n{bx}\\ntarget shape {by.shape}\\n{by}\")\n",
    "for xx in range(batch_size):\n",
    "    for t in range(block_size):\n",
    "        print(f\"input {bx[xx][:t+1]} have target {by[xx][t]}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss: 4.701944828033447\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "torch.Size([1, 1])\n",
      "generate test:\n",
      "S3fh$-M$gCjxvbRj;pGGju;TgCjXOca!CVtTbV$JSV;xZ$Q!U-Q?3faeDvrVHCDq-mc;ai?Oyvh&ymnk&yhsEXNC&yeAUu'Q?Ifn\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "torch.manual_seed(1337)\n",
    "\n",
    "# Bigram\n",
    "class BigramLanuageModel(nn.Module):\n",
    "    def __init__(self, vocab_size) -> None:\n",
    "        super().__init__()\n",
    "        # Embedding 层用来初始化每个char下一个char的概率，其原本用于词嵌入，这里可以用来方便的批量提取概率\n",
    "        # 在Bigram中我们的需求是给出idx得到该词的概率，这与Embedding层的作用相同，遂用此层\n",
    "        self.idx_to_probs = nn.Embedding(vocab_size, vocab_size)\n",
    "    \n",
    "    def forward(self, batch_x, batch_y=None):\n",
    "        logits = self.idx_to_probs(batch_x)\n",
    "        # 这里我们的batch_y 是（B, T）而cross_entropy 需要（B, C, T）或者是说（B, C）\n",
    "        # 这边相当于把一批的多个词汇给转化成了批次\n",
    "        if batch_y is None:\n",
    "            loss = None\n",
    "        else:\n",
    "            B, T, C  = logits.shape\n",
    "            loss = F.cross_entropy(logits.view(B*T, C), batch_y.view(B*T))\n",
    "        return logits, loss\n",
    "    \n",
    "    def generater(self, batch_x, max_size):\n",
    "        # 这里的batch_x 是(batch, 1)因为仅凭一个词汇预测 \n",
    "        for _ in range(max_size):\n",
    "            logits, _ = self.forward(batch_x[:, -1])\n",
    "            probs = F.softmax(logits, dim=1)\n",
    "            batch_x_next = torch.multinomial(probs, 1, replacement=True)\n",
    "            batch_x = torch.cat((batch_x, batch_x_next), dim=1)\n",
    "        return batch_x\n",
    "        \n",
    "m = BigramLanuageModel(vocabs_len)\n",
    "test_x = torch.zeros((1,1), dtype=torch.long)\n",
    "# —————————————————————————————————————————————————————— #\n",
    "print(f\"loss: {(m(bx, by))[1]}\")\n",
    "print(f\"generate test:{decoder(m.generater(test_x, 100)[0].tolist())}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "optim = torch.optim.AdamW(m.parameters(), lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.5986359119415283\n"
     ]
    }
   ],
   "source": [
    "batch_size = 32\n",
    "for _ in range(10000):\n",
    "    bx, by = get_batch()\n",
    "    logits, loss = m(bx, by)\n",
    "    optim.zero_grad(set_to_none=True)\n",
    "    loss.backward()\n",
    "    optim.step()\n",
    "print(loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "generate test:\n",
      "\n",
      "S:\n",
      "W'llst,\n",
      "Ung eto d by. d byo irex\n",
      "I y. dig the.\n",
      "CHes fepee ft thofe t RAMAMAbe t;'t'len ll on y\n",
      "Iforalke.\n",
      "TY hineeg.\n",
      "No wo d atas bal\n",
      "Thenor r s tow ceen, hoo fal thas, whilleco akHENOFor, Leatssen\n"
     ]
    }
   ],
   "source": [
    "print(f\"generate test:{decoder(m.generater(test_x, 200)[0].tolist())}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3.7.2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
